#!/usr/bin/env python3
"""基于 TUN 设备的真实 VPN 服务器实现。

本模块实现了一个完整的 VPN 服务器，包含以下核心功能：
1. TUN 虚拟网卡创建与管理
2. 客户端连接管理与 IP 地址分配
3. 加密隧道建立（TLS）
4. IP 数据包转发与 NAT
5. 路由配置与流量控制

注意：需要 root 权限运行，仅支持 Linux 平台。
"""

import argparse
import asyncio
import fcntl
import json
import logging
import os
import ssl
import struct
import subprocess
from dataclasses import dataclass
from ipaddress import IPv4Address, IPv4Network
from typing import Dict, Optional

# TUN 设备相关常量
TUNSETIFF = 0x400454ca  # ioctl 命令：设置 TUN/TAP 接口
IFF_TUN = 0x0001        # TUN 设备（三层，IP 包）
IFF_NO_PI = 0x1000      # 不包含额外的包信息头


@dataclass
class ClientSession:
    """VPN 客户端会话信息。"""
    
    client_id: str                      # 客户端唯一标识
    virtual_ip: IPv4Address             # 分配给客户端的虚拟 IP
    reader: asyncio.StreamReader        # TLS 连接读端
    writer: asyncio.StreamWriter        # TLS 连接写端
    last_active: float                  # 最后活跃时间


@dataclass
class ServerConfig:
    """VPN 服务器配置。"""
    
    bind_host: str                      # 监听地址
    bind_port: int                      # 监听端口
    auth_token: str                     # 认证令牌
    certfile: str                       # 服务器证书
    keyfile: str                        # 服务器私钥
    tun_name: str                       # TUN 设备名称
    vpn_network: IPv4Network            # VPN 虚拟网段
    server_ip: IPv4Address              # 服务器在 VPN 中的 IP
    physical_interface: str             # 物理网卡名称（用于 NAT）
    enable_nat: bool                    # 是否启用 NAT


class TunDevice:
    """TUN 虚拟网卡管理器。"""
    
    def __init__(self, name: str):
        self.name = name
        self.fd: Optional[int] = None
    
    def create(self) -> int:
        """创建 TUN 设备并返回文件描述符。"""
        try:
            # 打开 TUN 设备
            self.fd = os.open("/dev/net/tun", os.O_RDWR)
            
            # 配置为 TUN 模式（三层 IP）
            ifr = struct.pack("16sH", self.name.encode()[:15], IFF_TUN | IFF_NO_PI)
            fcntl.ioctl(self.fd, TUNSETIFF, ifr)
            
            logging.info(f"TUN 设备 {self.name} 创建成功")
            return self.fd
        except Exception as e:
            logging.error(f"创建 TUN 设备失败: {e}")
            raise
    
    def configure(self, ip_addr: str, netmask: str = "255.255.255.0"):
        """配置 TUN 设备的 IP 地址。"""
        try:
            # 设置 IP 地址
            subprocess.run(
                ["ip", "addr", "add", f"{ip_addr}/{netmask}", "dev", self.name],
                check=True
            )
            
            # 启动接口
            subprocess.run(
                ["ip", "link", "set", self.name, "up"],
                check=True
            )
            
            logging.info(f"TUN 设备 {self.name} 配置完成: {ip_addr}/{netmask}")
        except subprocess.CalledProcessError as e:
            logging.error(f"配置 TUN 设备失败: {e}")
            raise
    
    def read(self, size: int = 2048) -> bytes:
        """从 TUN 设备读取 IP 数据包。"""
        if self.fd is None:
            raise RuntimeError("TUN 设备未创建")
        return os.read(self.fd, size)
    
    def write(self, data: bytes) -> int:
        """向 TUN 设备写入 IP 数据包。"""
        if self.fd is None:
            raise RuntimeError("TUN 设备未创建")
        return os.write(self.fd, data)
    
    def close(self):
        """关闭 TUN 设备。"""
        if self.fd is not None:
            os.close(self.fd)
            self.fd = None
            logging.info(f"TUN 设备 {self.name} 已关闭")


class VPNServer:
    """VPN 服务器主类。"""
    
    def __init__(self, config: ServerConfig):
        self.config = config
        self.tun = TunDevice(config.tun_name)
        self.clients: Dict[str, ClientSession] = {}
        self.ip_pool = list(config.vpn_network.hosts())[1:]  # 跳过服务器 IP
        self.allocated_ips: set = set()
    
    def allocate_ip(self) -> Optional[IPv4Address]:
        """为新客户端分配虚拟 IP。"""
        for ip in self.ip_pool:
            if ip not in self.allocated_ips and ip != self.config.server_ip:
                self.allocated_ips.add(ip)
                return ip
        return None
    
    def release_ip(self, ip: IPv4Address):
        """释放客户端 IP。"""
        self.allocated_ips.discard(ip)
    
    def setup_nat(self):
        """配置 NAT 规则，允许 VPN 客户端访问外网。"""
        if not self.config.enable_nat:
            return
        
        try:
            # 启用 IP 转发
            subprocess.run(
                ["sysctl", "-w", "net.ipv4.ip_forward=1"],
                check=True
            )
            
            # 配置 iptables NAT（MASQUERADE）
            subprocess.run([
                "iptables", "-t", "nat", "-A", "POSTROUTING",
                "-s", str(self.config.vpn_network),
                "-o", self.config.physical_interface,
                "-j", "MASQUERADE"
            ], check=True)
            
            # 允许转发
            subprocess.run([
                "iptables", "-A", "FORWARD",
                "-i", self.config.tun_name,
                "-j", "ACCEPT"
            ], check=True)
            
            subprocess.run([
                "iptables", "-A", "FORWARD",
                "-o", self.config.tun_name,
                "-j", "ACCEPT"
            ], check=True)
            
            logging.info("NAT 规则配置成功")
        except subprocess.CalledProcessError as e:
            logging.error(f"配置 NAT 失败: {e}")
            raise
    
    def cleanup_nat(self):
        """清理 NAT 规则。"""
        if not self.config.enable_nat:
            return
        
        try:
            subprocess.run([
                "iptables", "-t", "nat", "-D", "POSTROUTING",
                "-s", str(self.config.vpn_network),
                "-o", self.config.physical_interface,
                "-j", "MASQUERADE"
            ], stderr=subprocess.DEVNULL)
            
            subprocess.run([
                "iptables", "-D", "FORWARD",
                "-i", self.config.tun_name,
                "-j", "ACCEPT"
            ], stderr=subprocess.DEVNULL)
            
            subprocess.run([
                "iptables", "-D", "FORWARD",
                "-o", self.config.tun_name,
                "-j", "ACCEPT"
            ], stderr=subprocess.DEVNULL)
            
            logging.info("NAT 规则已清理")
        except Exception as e:
            logging.warning(f"清理 NAT 规则时出错: {e}")
    
    async def handle_client_control(
        self,
        reader: asyncio.StreamReader,
        writer: asyncio.StreamWriter
    ):
        """处理客户端控制连接（认证与配置）。"""
        peer = writer.get_extra_info("peername")
        logging.info(f"新客户端连接: {peer}")
        
        try:
            # 读取认证信息
            auth_data = await asyncio.wait_for(reader.readline(), timeout=10)
            auth_msg = json.loads(auth_data.decode("utf-8"))
            
            if auth_msg.get("token") != self.config.auth_token:
                raise PermissionError("认证失败")
            
            client_id = auth_msg.get("client_id", str(peer))
            
            # 分配虚拟 IP
            virtual_ip = self.allocate_ip()
            if virtual_ip is None:
                raise RuntimeError("IP 地址池已耗尽")
            
            # 创建会话
            session = ClientSession(
                client_id=client_id,
                virtual_ip=virtual_ip,
                reader=reader,
                writer=writer,
                last_active=asyncio.get_event_loop().time()
            )
            self.clients[client_id] = session
            
            # 发送配置信息
            config_msg = {
                "status": "ok",
                "virtual_ip": str(virtual_ip),
                "server_ip": str(self.config.server_ip),
                "netmask": str(self.config.vpn_network.netmask),
                "network": str(self.config.vpn_network)
            }
            writer.write(json.dumps(config_msg).encode("utf-8") + b"\n")
            await writer.drain()
            
            logging.info(f"客户端 {client_id} 已分配 IP: {virtual_ip}")
            
            # 处理隧道数据
            await self.handle_tunnel_data(session)
            
        except Exception as e:
            logging.warning(f"客户端 {peer} 处理失败: {e}")
        finally:
            # 清理会话
            if client_id in self.clients:
                session = self.clients.pop(client_id)
                self.release_ip(session.virtual_ip)
                logging.info(f"客户端 {client_id} 已断开")
            
            try:
                writer.close()
                await writer.wait_closed()
            except Exception:
                pass
    
    async def handle_tunnel_data(self, session: ClientSession):
        """处理隧道数据双向转发。"""
        # 启动两个任务：从客户端读取 -> TUN，从 TUN 读取 -> 客户端
        await asyncio.gather(
            self.client_to_tun(session),
            self.tun_to_client(session),
            return_exceptions=True
        )
    
    async def client_to_tun(self, session: ClientSession):
        """从客户端读取数据包并写入 TUN 设备。"""
        try:
            while True:
                # 读取包长度（4 字节）
                length_data = await session.reader.readexactly(4)
                packet_length = struct.unpack("!I", length_data)[0]
                
                # 读取 IP 数据包
                packet = await session.reader.readexactly(packet_length)
                
                # 写入 TUN 设备
                loop = asyncio.get_event_loop()
                await loop.run_in_executor(None, self.tun.write, packet)
                
                session.last_active = loop.time()
                
        except asyncio.IncompleteReadError:
            logging.debug(f"客户端 {session.client_id} 连接关闭")
        except Exception as e:
            logging.error(f"client_to_tun 错误: {e}")
    
    async def tun_to_client(self, session: ClientSession):
        """从 TUN 设备读取数据包并发送给对应客户端。"""
        loop = asyncio.get_event_loop()
        
        try:
            while True:
                # 从 TUN 读取 IP 包
                packet = await loop.run_in_executor(None, self.tun.read, 2048)
                
                if len(packet) < 20:  # IP 头最小 20 字节
                    continue
                
                # 解析目标 IP（IP 头偏移 16-19 字节）
                dst_ip = IPv4Address(packet[16:20])
                
                # 查找目标客户端
                target_session = None
                for client in self.clients.values():
                    if client.virtual_ip == dst_ip:
                        target_session = client
                        break
                
                if target_session is None:
                    continue
                
                # 发送给客户端（长度 + 数据）
                length_data = struct.pack("!I", len(packet))
                target_session.writer.write(length_data + packet)
                await target_session.writer.drain()
                
        except Exception as e:
            logging.error(f"tun_to_client 错误: {e}")
    
    async def tun_reader_loop(self):
        """持续从 TUN 读取并分发数据包。"""
        loop = asyncio.get_event_loop()
        
        while True:
            try:
                packet = await loop.run_in_executor(None, self.tun.read, 2048)
                
                if len(packet) < 20:
                    continue
                
                # 解析目标 IP
                dst_ip = IPv4Address(packet[16:20])
                
                # 如果目标是 VPN 网段内的客户端，转发给对应客户端
                if dst_ip in self.config.vpn_network:
                    for client in self.clients.values():
                        if client.virtual_ip == dst_ip:
                            length_data = struct.pack("!I", len(packet))
                            client.writer.write(length_data + packet)
                            await client.writer.drain()
                            break
                
            except Exception as e:
                logging.error(f"TUN 读取循环错误: {e}")
                await asyncio.sleep(0.1)
    
    async def start(self):
        """启动 VPN 服务器。"""
        # 创建并配置 TUN 设备
        self.tun.create()
        self.tun.configure(
            str(self.config.server_ip),
            str(self.config.vpn_network.prefixlen)
        )
        
        # 配置 NAT
        self.setup_nat()
        
        # 创建 TLS 上下文
        ssl_ctx = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
        ssl_ctx.load_cert_chain(self.config.certfile, self.config.keyfile)
        
        # 启动 TCP 服务器
        server = await asyncio.start_server(
            self.handle_client_control,
            self.config.bind_host,
            self.config.bind_port,
            ssl=ssl_ctx
        )
        
        logging.info(f"VPN 服务器启动: {self.config.bind_host}:{self.config.bind_port}")
        logging.info(f"VPN 网段: {self.config.vpn_network}")
        
        # 启动 TUN 读取循环
        tun_task = asyncio.create_task(self.tun_reader_loop())
        
        try:
            async with server:
                await server.serve_forever()
        finally:
            tun_task.cancel()
            self.cleanup_nat()
            self.tun.close()


def parse_args() -> ServerConfig:
    parser = argparse.ArgumentParser(description="基于 TUN 的真实 VPN 服务器")
    parser.add_argument("--bind-host", default="0.0.0.0", help="监听地址")
    parser.add_argument("--bind-port", type=int, default=8443, help="监听端口")
    parser.add_argument("--auth-token", required=True, help="认证令牌")
    parser.add_argument("--cert", required=True, help="服务器证书路径")
    parser.add_argument("--key", required=True, help="服务器私钥路径")
    parser.add_argument("--tun-name", default="vpn0", help="TUN 设备名称")
    parser.add_argument("--vpn-network", default="10.8.0.0/24", help="VPN 虚拟网段")
    parser.add_argument("--server-ip", default="10.8.0.1", help="服务器虚拟 IP")
    parser.add_argument("--physical-interface", default="eth0", help="物理网卡名称")
    parser.add_argument("--enable-nat", action="store_true", help="启用 NAT")
    parser.add_argument("--log-level", default="INFO", help="日志等级")
    
    args = parser.parse_args()
    
    logging.basicConfig(
        level=getattr(logging, args.log_level.upper()),
        format="%(asctime)s | %(levelname)s | %(message)s"
    )
    
    return ServerConfig(
        bind_host=args.bind_host,
        bind_port=args.bind_port,
        auth_token=args.auth_token,
        certfile=args.cert,
        keyfile=args.key,
        tun_name=args.tun_name,
        vpn_network=IPv4Network(args.vpn_network),
        server_ip=IPv4Address(args.server_ip),
        physical_interface=args.physical_interface,
        enable_nat=args.enable_nat
    )


def main():
    if os.geteuid() != 0:
        print("错误：需要 root 权限运行")
        return 1
    
    config = parse_args()
    server = VPNServer(config)
    
    try:
        asyncio.run(server.start())
    except KeyboardInterrupt:
        logging.info("服务器已停止")
    
    return 0


if __name__ == "__main__":
    exit(main())
